# coding:utf-8
from ..op_basic import as_op
try:
    from sklearn.cross_validation import KFold
    __USE_OLD_KOLD = True
except Exception:
    from sklearn.model_selection import KFold
    __USE_OLD_KOLD = False

import logging


def get_train_test(X, y, kfold=False, kfold_k=None, kfold_shuffle=False):
    if kfold is False:
        yield X, y, None, None
    else:
        if __USE_OLD_KOLD:
            kf = KFold(X.shape[0], n_folds=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf:
                yield X[train_index], y[train_index], X[test_index], y[test_index]
        else:
            kf = KFold(n_splits=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf.split(X):
                yield X[train_index], y[train_index], X[test_index], y[test_index]

def get_train_test_index(X, y, kfold=False, kfold_k=None, kfold_shuffle=False):
    if kfold is False:
        yield X, y, None, None
    else:
        if __USE_OLD_KOLD:
            kf = KFold(X.shape[0], n_folds=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf:
                yield train_index, test_index
        else:
            kf = KFold(n_splits=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf.split(X):
                yield train_index, test_index

@as_op('train_fold')
def train_fold(y, df, use_model, model_para, scorer=None, kfold_k=4, kfold_shuffle=False):
    logging.info('Use %d Folds...' % kfold_k)

    all_score = 0

    y_df = df[y]
    X_df = df.drop([y], axis=1)
    X = X_df.values
    y = y_df.values

    for i, (X_train, y_train, X_test, y_test) in enumerate(get_train_test(X, y, True, kfold_k, kfold_shuffle)):
        # X_train, X_test = X[train_index], X[test_index]
        # y_train, y_test = y[train_index], y[test_index]
        model = use_model(**model_para)
        model.fit(X_train, y_train)
        if scorer is not None:
            score = scorer(model, X_test, y_test)
        else:
            score = model.score(X_test, y_test)
        all_score += score
        logging.info('Fold %d/%d Score: %f ' % (i + 1, kfold_k, score))
    logging.info('Avg Score: %f ' % (all_score / kfold_k))
